import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import json

# 加载训练时保存的模型
model = load_model('best_model.h5')

# 加载 class_indices (在训练时需要保存这个文件)
class_indices = {
    "glioma": 0,
    "meningioma": 1,
    "notumor": 2,
    "pituitary": 3
}

# 将 class_indices 的键值对翻转
class_indices = {v: k for k, v in class_indices.items()}

# 预测函数
def predict_image(img_path):
    img = image.load_img(img_path, target_size=(150, 150))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0) / 255.0

    prediction = model.predict(img_array)
    predicted_class = class_indices[np.argmax(prediction)]
    
    return predicted_class

# 示例预测
try:
    if path != "@END@":
        img_path = '../archive/Testing/meningioma/Te-me_0145.jpg'  # 替换成你想预测的图像路径

        probabilities = predict_image(img_path)
        predicted_class = class_indices[np.argmax(probabilities)]
        cancer_name = predicted_class
        probability_str = str(max(probabilities[0]))

        # print(predicted_class)

        # 读取并且展示图像
        img = mpimg.imread(img_path)
        plt.figure()
        plt.imshow(img)
        plt.axis('off')
        # 添加预测肿瘤的结果以及对应预测的概率于图像上
        plt.text(0.5, 0.95, cancer_name + ", probability:" + probability_str,
                horizontalalignment='center', verticalalignment='center', transform=plt.gca().transAxes,
                fontsize=12, color='white', bbox=dict(facecolor='black', alpha=0.5))
        plt.show()
except Exception as e:
    print("输入格式错误或发生异常：", e)
